## covid (+IAV & HC) PBMC
library(Seurat)
library(tidyseurat)
library(tidyverse)
library(ggpubr)
library(pheatmap)
library(RColorBrewer)
library(greekLetters)
source('00_util_scripts/mod_seurat.R', echo = FALSE)

# define markers ----------
tcr <- c('CD3D','CD3E','CD3G','CD247','TRAC','TRBC1','TRBC2','TRDC','TRGC1','TRGC2')

bcr <- c('CD79A','CD79B','IGKC','IGLC2','IGLC3','IGLC6','IGLC7','IGHA1','IGHA2','IGHG1','IGHG2','IGHG3','IGHG4','IGHD','IGHE','IGHM')

Ahmed <- c('CD19', 'CD5', 'TRAC', 'TRBC1', 'TRBC2', 'IGHD')

all_marker <- unique(c(bcr, tcr, Ahmed))

write_rds(all_marker, 'DE_cells/ref/de_markers.rds')

# covid-19 & IAV data -------------
blish_sobj <- read_rds('DE_cells/data/blish_cleaned.rds')
bgi_sobj <- read_rds('DE_cells/data/bgi_soupx.rds')

monaco <- celldex::MonacoImmuneData()

blish_sobj <- blish_sobj |> filter(nFeature_RNA >= 200 & percent.mt < 10)

blish_sobj <- blish_sobj |> quick_process_seurat()

blish_sobj <- blish_sobj |>
  mark_cell_type_singler(ref = monaco, fine_label = TRUE, new_label = 'monaco_label')

blish_sobj |>
  DietSeurat() |>
  write_rds('DE_cells/data/blish_cleaned.rds')

get_abundance_sc_wide(blish_sobj, all_marker) %>%
  right_join(as_tibble(blish_sobj)) %>%
  select(!matches("PC_|harmony")) %>%
  data.table::fwrite(str_glue("DE_cells/results/blish_de_meta.csv.gz"))

bgi_sobj <- read_rds('DE_cells/data/BGI_cleaned.rds')

bgi_sobj <- bgi_sobj |> filter(nFeature_RNA >= 200 & percent.mt < 10)

bgi_sobj <- quick_process_seurat(bgi_sobj)

bgi_sobj <- bgi_sobj |>
  mark_cell_type_singler(ref = monaco, fine_label = TRUE, new_label = 'monaco_label')

bgi_sobj |>
  DietSeurat() |>
  write_rds('DE_cells/data/BGI_cleaned.rds')

get_abundance_sc_wide(bgi_sobj, all_marker) %>%
  right_join(as_tibble(bgi_sobj)) %>%
  select(!matches("PC_|harmony")) %>%
  data.table::fwrite(str_glue("DE_cells/results/bgi_de_meta.csv.gz"))

plot_de_cell_filtering <- function(sobj) {
  meta_data <- sobj@meta.data %>%
    as_tibble(rownames = 'rowname')
  
  # note: Blish have no IGLC6 + TRBC1
  if (none(rownames(sobj), \(x)x == 'IGLC6')) {
    meta_data <- meta_data %>%
      rename_with(~ 'Stage', starts_with('Status'))
  }
  
  fill_expr_mat <- function(meta_df){
    if (every(all_marker, \(x)x %in% colnames(meta_df))) meta_df
    
    unexist_marker <- all_marker %>%
      discard(\(x)x %in% colnames(meta_df))
    
    matrix(nrow = dim(meta_df)[[1]], ncol = length(unexist_marker)) %>%
      as.data.frame() %>%
      set_names(unexist_marker) %>%
      bind_cols(meta_df)
  }

  marker_counts <- FetchData(sobj, all_marker) %>%
    as_tibble(rownames = 'rowname') %>%
    right_join(meta_data) %>%
    filter(scrublet_call == 'Singlet' & nFeature_RNA >= 200 & percent.mt < 10) %>%
    mutate(IGL = IGKC| IGLC2 | IGLC3 | IGLC6 |IGLC7,
           IGH = IGHA1 | IGHA2 | IGHG1 | IGHG2 | IGHG3 | IGHG4 | IGHD | IGHE | IGHM,
           TRG = TRGC1 | TRGC2,
           TRB = TRBC1 | TRBC2,
           TCRab = TRAC & TRB,
           TCRgd = TRG & TRDC,
           TCR = TCRab | TCRgd,
           CD3 = CD3D & CD3E & CD3G & CD247)
  
  cell_num_table <- marker_counts %>%
    group_by(Stage) %>%
    reframe(Total_cells = n(),
            DE_cells = sum(CD19 & CD5 & TRAC & TRB)) %>%
    mutate(DE_fraction = DE_cells / Total_cells) %>%
    column_to_rownames('Stage') %>%
    t() %>%
    as.data.frame() %>%
    rownames_to_column('Cell type') %>%
    bind_rows()
  
  cell_num_table <- meta_data %>%
    dplyr::count(Stage, singleR) %>%
    pivot_wider(names_from = 'Stage', values_from = 'n') %>%
    rename_with(~'Cell type', matches('singler')) %>%
    bind_rows(cell_num_table)
  
  # plot heatmap in DE cell filtering --------
  ahmed_count <- marker_counts %>%
    filter(Stage %in% c('Ctrl', 'Healthy')) %>%
    group_by(singleR) %>%
    reframe(sum = n(),
            de1 = sum(CD19 > 0),
            de2 = sum(CD19 & CD5),
            de3 = sum(CD19 & CD5 & TRAC),
            de4 = sum(CD19 & CD5 & TRAC & TRB),
            de5 = sum(CD19 & CD5 & TRAC & TRB & IGHD))
  
  ahmed_mat <- ahmed_count %>%
    mutate(across(where(is.numeric), ~./sum)) %>%
    select(-sum) %>%
    column_to_rownames('singleR') %>%
    t()
  
  a <- 'CD19+'
  b <- 'CD5+'
  c <- str_glue('TCR', greek$alpha, '+')
  d <- str_glue('TCR', greek$alpha, greek$beta, '+')
  e <- 'IgD+'
  rownames(ahmed_mat) <- c(a,
                           str_glue(a, b),
                           str_glue(a, b, c),
                           str_glue(a, b, d),
                           str_glue(a, b, d, e))
  
  pheatmap(ahmed_mat,
           color = colorRampPalette(c('white','coral'))(30),
           scale = 'row',
           cluster_rows = FALSE,
           cluster_cols = FALSE,
           display_numbers = formattable::percent(ahmed_mat),
           number_color = 'black') %>%
    ggplotify::as.ggplot()
  
  ahmed_heatmap <- last_plot()
  
  # Zhang XK's criteria ---------
  xk_count <- marker_counts %>%
    filter(Stage %in% c('Ctrl', 'Healthy')) %>%
    group_by(singleR) %>%
    reframe(sum = n(),
            de1 = sum(CD79A & CD79B),
            de2 = sum(CD79A & CD79B & IGL),
            de3 = sum(CD79A & CD79B & IGL & IGH),
            de4 = sum(CD79A & CD79B & IGL & IGH & CD3),
            de5 = sum(CD79A & CD79B & IGL & IGH & CD3 & TCR))
  
  xk_mat <- xk_count %>%
    mutate(across(where(is.numeric), ~./sum)) %>%
    select(-sum) %>%
    column_to_rownames('singleR') %>%
    t()
  
  rownames(xk_mat) <- c('CD79+', 'CD79+IGL+', 'CD79+IGL+IGH+', 'CD79+IGL+IGH+CD3+', 'CD79+IGL+IGH+CD3+TCR+')
  
  pheatmap(xk_mat,
           color = colorRampPalette(c('white','coral'))(30),
           scale = 'row',
           cluster_rows = FALSE,
           cluster_cols = FALSE,
           number_format = '%.2f',
           display_numbers = formattable::percent(xk_mat),
           number_color = 'black') %>%
    ggplotify::as.ggplot()
  
  xk_heatmap <- last_plot()
  
  list(ahmed_mat, 
       ahmed_heatmap,
       xk_mat,
       xk_heatmap,
       cell_num_table)
}

# invoke func --------
de_blish_res <- blish_sobj %>%
  plot_de_cell_filtering()

de_blish_res[[5]] %>% flextable::flextable()

de_bgi_res <- bgi_sobj %>%
  plot_de_cell_filtering()

cairo_pdf("DE_cells/figures/Sfig1_de_filtering_adj.pdf",
          width = 10,
          height = 7.8)

de_bgi_res[[2]] / de_bgi_res[[4]] +
  patchwork::plot_annotation(tag_levels = 'A') &
  theme(plot.tag = element_text(size = 32))

dev.off()

ggsave("DE_cells/figures/Sfig1_de_filtering_adj.png",
       width = 10,
       height = 7.8,
       units = 'in')

write.csv(de_bgi_res[[1]], 'DE_cells/results/Sfig1_de_filtering_ahmed_adj.csv')
write.csv(de_bgi_res[[3]], 'DE_cells/results/Sfig1_de_filtering_zxk_adj.csv')

write.csv(de_bgi_res[[5]], 'DE_cells/results/table1_bgi_cell_num_adj.csv')
write.csv(de_blish_res[[5]], 'DE_cells/results/table1_blish_cell_num.csv')

